// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/exec/enforce_single_row.h>

namespace kumo::pollux::exec {
    EnforceSingleRow::EnforceSingleRow(
        int32_t operatorId,
        DriverCtx *driverCtx,
        const std::shared_ptr<const core::EnforceSingleRowNode> &planNode)
        : Operator(
            driverCtx,
            planNode->outputType(),
            operatorId,
            planNode->id(),
            "EnforceSingleRow") {
        isIdentityProjection_ = true;
    }

    void EnforceSingleRow::addInput(RowVectorPtr input) {
        auto numInput = input->size();
        POLLUX_CHECK_NE(
            numInput, 0, "EnforceSingleRow::addInput received empty set of rows");
        if (input_ == nullptr) {
            POLLUX_USER_CHECK_EQ(
                numInput,
                1,
                "Expected single row of input. Received {} rows.",
                numInput);
            input_ = std::move(input);
        } else {
            POLLUX_USER_FAIL(
                "Expected single row of input. Received {} extra rows.", numInput);
        }
    }

    RowVectorPtr EnforceSingleRow::getOutput() {
        if (!noMoreInput_) {
            return nullptr;
        }

        return std::move(input_);
    }

    void EnforceSingleRow::noMoreInput() {
        if (!noMoreInput_ && input_ == nullptr) {
            // We have not seen any data. Return a single row of all nulls.
            input_ = BaseVector::create<RowVector>(outputType_, 1, pool());
            for (auto &child: input_->children()) {
                child->resize(1);
                child->set_null(0, true);
            }
        }
        Operator::noMoreInput();
    }

    bool EnforceSingleRow::isFinished() {
        return noMoreInput_ && input_ == nullptr;
    }
} // namespace kumo::pollux::exec
